# DxdDataV1测试类
import unittest
import numpy as np
import matplotlib.pyplot as plt
import torch
from apps.dhlp.dxd_data_v1 import DxdDataV1

class TDxdDataV1(unittest.TestCase):
    # python -m unittest -v uts.apps.dhlp.t_dxd_data_v1.TDxdDataV1.test_get_rid_recs
    def test_get_rid_recs(self):
        rec_db, dev_recs, raw_datas = DxdDataV1.load_raw_datas(params={})
        # 1. 找到记录数最多的设备
        did = DxdDataV1.get_max_recs_dev(dev_recs=dev_recs)
        print(f'did: {did};')
        recs = DxdDataV1.get_rid_recs(did, dev_recs, raw_datas)
        # {'rid': {'data': [1.0, 2.0, ..., 1.0],'label': 0 # 0-代表正常; 1-代表雷电;}}
        for k, v in recs.items():
            print(f'{k}: {len(v["data"])}+{v["label"]};')

    # python -m unittest -v uts.apps.dhlp.t_dxd_data_v1.TDxdDataV1.test_extract_frames
    def test_extract_frames(self):
        rec_db, dev_recs, raw_datas = DxdDataV1.load_raw_datas(params={})
        # 1. 找到记录数最多的设备
        did = DxdDataV1.get_max_recs_dev(dev_recs=dev_recs)
        print(f'did: {did};')
        tl_recs = DxdDataV1.get_rid_recs(did, dev_recs, raw_datas)
        # {'rid': {'data': [1.0, 2.0, ..., 1.0],'label': 0 # 0-代表正常; 1-代表雷电;}}
        for rid, v in tl_recs.items():
            print(f'rid: {rid};')
            break
        frames = DxdDataV1.extract_frames(tl_recs[rid]['data'], tl_recs[rid]['label'])
        print(f'frames: {frames.shape};')

        # 50~59：距离雷电发生1~10分钟的帧
        # 40~49：距离雷电发生11~20分钟
        # 30~39：距离雷电发生21~30分钟

    # python -m unittest -v uts.apps.dhlp.t_dxd_data_v1.TDxdDataV1.test_geneate_norm_recs_dict
    def test_geneate_norm_recs_dict(self):
        rec_db, dev_recs, raw_datas = DxdDataV1.load_raw_datas(params={})
        # 1. 找到记录数最多的设备
        did = DxdDataV1.get_max_recs_dev(dev_recs=dev_recs)
        print(f'did: {did};')
        # tl_recs {'rid': {'data': [1.0, 2.0, ..., 1.0],'label': 0 # 0-代表正常; 1-代表雷电;}}
        recs = DxdDataV1.get_rid_recs(did, dev_recs, raw_datas)
        norm_recs = DxdDataV1.generate_norm_recs_dict(recs)
        for k, v in norm_recs.items():
            recs[k] = v
        for k, v in recs.items():
            print(f'### {k}: {len(v["data"])}, {v["label"]};')

    # python -m unittest -v uts.apps.dhlp.t_dxd_data_v1.TDxdDataV1.test_generate_fn
    def test_generate_fn(self):
        seq = 2008
        fn = DxdDataV1.generate_fn(seq)
        print(f'@@@ fn: {fn};')

    # python -m unittest -v uts.apps.dhlp.t_dxd_data_v1.TDxdDataV1.test_save_frames
    def test_save_frames(self):
        rec_db, dev_recs, raw_datas = DxdDataV1.load_raw_datas(params={})
        # 1. 找到记录数最多的设备
        did = DxdDataV1.get_max_recs_dev(dev_recs=dev_recs)
        print(f'did: {did};')
        # tl_recs {'rid': {'data': [1.0, 2.0, ..., 1.0],'label': 0 # 0-代表正常; 1-代表雷电;}}
        recs = DxdDataV1.get_rid_recs(did, dev_recs, raw_datas)
        norm_recs = DxdDataV1.generate_norm_recs_dict(recs)
        for k, v in norm_recs.items():
            recs[k] = v
        seq = 0
        DxdDataV1.save_frames(seq, did, recs)

    # python -m unittest -v uts.apps.dhlp.t_dxd_data_v1.TDxdDataV1.test_load_ds_txt
    def test_load_ds_txt(self):
        ds_txt_fn = 'work/datasets/dhlp/v1/ds.txt'
        rids, rf0s, rf1s, rf2s, rf3s = DxdDataV1.load_ds_txt(ds_txt_fn)
        print(f'数据可视化过程')
        num = 0
        for k, v in rf1s.items():
            if k[0] != 'n':
                print(f'### {k}: {len(v)};')
                # 读入原始数据
                cnt = len(v)
                for i in range(cnt):
                    X_pt = torch.load(v[i]['fn'], weights_only=True)
                    X_np = X_pt.numpy()
                    X_np = 20 * np.log10(X_np + 51.0)
                    print(f'### X_np: {X_np.shape};')
                    x = np.arange(X_np.shape[0])
                    plt.plot(x, X_np)
                    plt.show()
                num += 1
                if num >= 3:
                    break

    # python -m unittest -v uts.apps.dhlp.t_dxd_data_v1.TDxdDataV1.test_show_red_alarm_frames
    def test_show_red_alarm_frames(self):
        ds_txt_fn = 'work/datasets/dhlp/v1/ds.txt'
        did = 'F56AC57A1181CD'
        rid = '16722'
        DxdDataV1.show_red_alarm_frames(ds_txt_fn, did, rid)

    # python -m unittest -v uts.apps.dhlp.t_dxd_data_v1.TDxdDataV1.test_show_red_alarm_frames002
    def test_show_red_alarm_frames002(self):
        ds_txt_fn = 'work/datasets/dhlp/v1/ds.txt'
        did = 'F56AC57A1181CD'
        rid = '16722'
        DxdDataV1.show_red_alarm_frames(ds_txt_fn, did, rid, mode=2)

    # python -m unittest -v uts.apps.dhlp.t_dxd_data_v1.TDxdDataV1.test_show_all_red_alarm_frames
    def test_show_all_red_alarm_frames(self):
        ds_txt_fn = 'work/datasets/dhlp/v1/ds.txt'
        DxdDataV1.show_all_red_alarm_frames(ds_txt_fn)

    # python -m unittest -v uts.apps.dhlp.t_dxd_data_v1.TDxdDataV1.test_show_all_red_alarm_frames002
    def test_show_all_red_alarm_frames002(self):
        ds_txt_fn = 'work/datasets/dhlp/v1/ds.txt'
        DxdDataV1.show_all_red_alarm_frames(ds_txt_fn, mode=2)

    def test_t1(self):
        data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]
        frame = data[3:6]
        X = torch.tensor(frame)
        frame1 = data[2:5]
        X = torch.vstack((X, torch.tensor(frame1)))